package com.we.flink.source;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.util.concurrent.ExecutorThreadFactory;

import com.we.flink.bean.FlinkCheckpointInfo;
import com.we.yarn.YarnClientUtils;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.exceptions.YarnException;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 自定义的 {@link YarnApplicationMonitorSource} Source 主要用来监听 Yarn 上运行的任务
 *
 * @author jasonlee
 */
public class YarnApplicationMonitorSource extends RichSourceFunction<FlinkCheckpointInfo> {
    private static final long serialVersionUID = 1L;
    private volatile boolean isRunning = true;
    private static List<ApplicationReport> flinkCheckpointInfos = new ArrayList<>(200);
    private static final SimpleDateFormat SDF = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    @Override
    public void open(Configuration parameters) throws Exception {
        // 定期扫描 Yarn 上任务的状态
        ScheduledExecutorService schedulerYarnService =
                Executors.newScheduledThreadPool(
                        1, new ExecutorThreadFactory("yarn-application-monitor-pool"));
        schedulerYarnService.scheduleWithFixedDelay(
                () -> {
                    synchronized (this) {
                        try {
                            flinkCheckpointInfos = YarnClientUtils.getApplicationReportList();
                        } catch (IOException | YarnException e) {
                            throw new RuntimeException(e);
                        }
                    }
                },
                60,
                60 * 5,
                TimeUnit.SECONDS);
    }

    @Override
    public void run(SourceContext<FlinkCheckpointInfo> ctx) {
        while (isRunning) {
            if (!flinkCheckpointInfos.isEmpty()) {
                flinkCheckpointInfos.forEach(
                        applicationReport -> {
                            FlinkCheckpointInfo flinkCheckpointInfo = new FlinkCheckpointInfo();
                            flinkCheckpointInfo.setApplicationId(
                                    String.valueOf(applicationReport.getApplicationId()));
                            flinkCheckpointInfo.setApplicationName(applicationReport.getName());
                            flinkCheckpointInfo.setStartTime(
                                    SDF.format(applicationReport.getStartTime()));
                            flinkCheckpointInfo.setTrackingUrl(applicationReport.getTrackingUrl());
                            ctx.collect(flinkCheckpointInfo);
                        });
                flinkCheckpointInfos.clear();
            }
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
    }
}
